package tests;

import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Scanner;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;

import stdp_rules.LogSTDP;
import utils.Metrics;
import utils.StateCamera;
import utils.StateCamera.CameraType;
import utils.SynapseRandomizer;
import utils.corpus_generator.SetTrainingPattern;
import utils.math_utils.Coordinate;
import utils.math_utils.MathUtils;
import utils.math_utils.RandStream;
import utils.math_utils.RandStream.Rand;
import utils.math_utils._Matrix;
import Jama.Matrix;

import components.LIFNeuron;
import components.Network;
import components.Synapse;

import connectors.GridConnect;

public class GrammarTest {
	public static final int numTrainSentences = 9000;
	public static final int numTestSentences = 1000;
	public static int trainEpochs;
	public static int testEpochs;
	public static final double timeStep = 1; //ms
	public static final double exposureISI = 30; //ms
	public static StateCamera camera;
	public static int numInputs = 576;
	public static int numRes = 1024;
	public static final int numOuts = 24;
	public static HashMap<Integer, LIFNeuron []> spikeGroups;
	public static int [][] sgIndexes;
	public static LIFNeuron[] inputs;
	public static LIFNeuron[] res;
	public static Synapse [][] rr;
	public static Synapse [][] ir;
	public static boolean [] trainingIn;
	public static boolean [] testingIn;
	public static double [][] teachTargs;
	public static double [][] testTargs;
	public static LogSTDP stdp;
	public static GridConnect gCon;
	static Scanner Keyboard = new Scanner(System.in);
	public static int numResThreads = 7;
	//public static int numInputThreads = 2;
	public static int numConsumerThreads = numResThreads + 1;
	private static double [] fireTimes = new double [numInputs];
	public static CountDownLatch latch;
	public static double [][] Corpus;
	public static Matrix wOut;
	public static double [] meanISIs;
	
	private static class Producer implements Runnable {
		private int next = 0;	
		private int [] corp;
		private StateCamera stateCamera;
		private boolean cameraOn;
		private final BlockingQueue<Runnable> queue;
		private LogSTDP stdp;
		private boolean stdpOn;
		private double time = 0.0;
		private boolean moreIns = true;
		private int epochs;
		Producer(BlockingQueue<Runnable> q, LogSTDP s, int [] corpus,
				StateCamera sc, int epochs) {
			queue = q;
			corp = corpus;
			spikeGroups = createGroups();
			stateCamera = sc;
			cameraOn = sc != null;
			stdp = s;
			stdpOn = s != null;
			this.epochs  = epochs;
			for(int i = 0; i < numInputs; i++){
				inputs[i].setReporter(false);
			}
			for(int i = 0; i < numRes; i++) {
				res[i].setReporter(cameraOn);
				res[i].setRecorder(stateCamera);
				
			}	
		}
		
		public void run () {
			int per = ((int)exposureISI * epochs)/100;
			int add = per;
			int i = 0;
			int c = 0;
			produce();
			while(moreIns) {
				if(queue.isEmpty()) {
					if(cameraOn){
						stateCamera.collectFrame();
					}
					if(i == add) {
						System.out.print(c +"%  ");
						add += per;
						c++;
					}
					produce();
					i++;
				}
			}
			System.out.println();
			latch.countDown();
			
		}	
		
		private void produce(){
			if(next >= epochs){
				moreIns = false;
				return;
			}
			try {	
				if(!queue.isEmpty()) {
					throw new InterruptedException();
				}
				if(time % exposureISI == 0) {
					for(int j = 0; j < 24; j++) {
						if(j != corp[next]){
							for(int k = 0; k < 24; k++)
								fireTimes[spikeGroups.get(j)[k].getIndex()] += exposureISI;	
						} else {
							double [] ft = new double [24];
							for(int k = 0; k < 24; k++) {
								ft[k] = fireTimes[sgIndexes[corp[next]][k]];
							}
							LIFNeuron[] sg = spikeGroups.get(corp[next]);
							queue.put(new runInputs(ft, sg,
									exposureISI/meanISIs[corp[next]],
										stdp, stdpOn, corp[next]));
						}
					}
					next++;
				}
				
				int index = 0;
				for(int i = 0; i < numResThreads; i++) {
					int size = (numResThreads - index) < numRes/numResThreads ?
							numResThreads - i : numRes/numResThreads;
					LIFNeuron [] p = new LIFNeuron[size];
					for(int j = 0; j < size; j++){
							p[j] = res[j + index];
					}
					queue.put(new runReservoir(p, stdp, stdpOn));
					index += size;
				}
				
				index = 0;
				
				time += timeStep;
				
			} catch (Exception e) {
				e.printStackTrace();
				System.err.println("ERROR: Tasks can only be added to empty queue.");
				System.exit(1);
			} 
			
		}
		
	}
	

	private static class Consumer implements Runnable {
		private final BlockingQueue<Runnable> queue;
		Consumer(BlockingQueue<Runnable> q) { queue = q; }
		
		public void run() {
			
			try {
				
				while(true) {
					queue.take().run();
				}
				
				
			} catch(InterruptedException e) {
				e.printStackTrace();
			}
			
			
		}
		
	}
	
	private static class runInputs implements Runnable {
		
		private final double [] fts;
		private final LIFNeuron [] firingInputNeurons;
		private final double meanPoissonInterval;
		private final LogSTDP stdp;
		private final boolean stdpOn;
		private final int index;
		
		public runInputs(double[] ft, LIFNeuron [] n, double mpi,
				LogSTDP s, boolean sOn, int ind) {
			fts = ft;
			firingInputNeurons = n;
			meanPoissonInterval = mpi;
			stdp = s;
			stdpOn = sOn;
			index = ind;
		}
		
		public void run(){
			for (int j = 0, n = fts.length; j < n; j++) {			
				if(fts[j] + firingInputNeurons[j].getLastSpike()
						<= firingInputNeurons[j].getTime()){
					firingInputNeurons[j].spike();
					if(stdpOn)
						stdp.applySTDPRule(firingInputNeurons[j]
								.getFanOut(), firingInputNeurons[j]
										.getTime()-timeStep, false);
					fts[j] = RandStream
							.nextPoissonInterval(meanPoissonInterval);
					try {
						fireTimes[sgIndexes[index][j]] = fts[j];
					}catch(ArrayIndexOutOfBoundsException e) {
						System.out.println("j: " + j);
						System.out.println("index " + index);
						System.out.println("sgIndexesR " + sgIndexes.length);
						if(sgIndexes.length > index){
							System.out.println("sgIndexesC " + sgIndexes[index].length);
						}
						System.exit(1);
					}
				}
				firingInputNeurons[j].update(timeStep);
			}
		}
		
	}
	
	private static class runReservoir implements Runnable {
		
		private final LIFNeuron [] firingResNeurons;
		private final LogSTDP stdp;
		private final boolean stdpOn;
		
		public runReservoir(LIFNeuron [] frn, LogSTDP s, boolean sOn) {
			firingResNeurons = frn;
			stdp  = s;
			stdpOn = sOn;
		}
		
		public void run() {
			for (int q = 0, n = firingResNeurons.length; q < n; q++) {
				firingResNeurons[q].update(timeStep);	
				if(stdpOn){
					stdp.applySTDPRule(firingResNeurons[q].getFanOut(),
							firingResNeurons[q].getTime()-timeStep, false);
					stdp.applySTDPRule(firingResNeurons[q].getFanIn(),
							firingResNeurons[q].getTime()-timeStep, true);
				}
			}
		}
		
	}
	
	private static class runForward implements Callable<Boolean> {
		
		private final int [] corp;
		private final int epochs;
		private final LogSTDP stdp;
		private final StateCamera camera;
		
		public runForward (int [] corp, int epochs, LogSTDP stdp,
				StateCamera camera) {
				this.corp = corp; 
				this.epochs = epochs;
				this.stdp = stdp;
				this.camera = camera;
			}
		
		public Boolean call(){
			BlockingQueue<Runnable> sharedQueue = new LinkedBlockingQueue<Runnable>();
			Thread producer = new Thread(new Producer(sharedQueue, stdp, corp, camera, epochs));
			Thread [] consumers = new Thread[numResThreads];
			
			for(int i = 0; i < numResThreads; i++) {
				consumers[i] = new Thread(new Consumer(sharedQueue));
			}
			producer.start();
			for(int i = 0; i < numResThreads; i++) {
				consumers[i].start();
			}
			
			return true;
		}
		
	}
	
	private static class solve implements Callable<Boolean> {
		private final double [][] Teach;
		private final StateCamera camera;
		public solve(double[][] Teach, StateCamera camera){
			this.Teach = Teach;
			this.camera = camera;
		}
		
		public Boolean call(){
			double [][] states = Arrays.copyOfRange(camera
					 .getCapturedStates(), 0, camera.getCapturedStates()
					 	.length-1);
			for(int i = 0; i < states.length; i++) {
				states[i] = MathUtils.concatArrays(states[i], Corpus[i]);
				for(int j = 0; j < states[i].length; j++){
					states[i][j] = states[i][j] + (0.0001 * (Math.random() - 0.5));
				}
			}
			 Matrix S = new Matrix(states);
			 Matrix D = new Matrix(Teach);
				System.out.println("States: " + states.length + " " + states[0].length);
				System.out.println("Teach: " + Teach.length + " " + Teach[0].length);
			 System.out.println("\nSolving...");
			 wOut = S.solve(D);	
			 latch.countDown();
			 return true;
		}
		
	}
	
	public static void main(String[] args) {
		String corpFile = "TCorpus.dat";
		double cos = 0.0;
		double MaxPred = 0.0;
		double VerbAgreement = 0.0;
		double [][] VerbD = new double[14][2];
		double [][] ROC = new double[2][200];
		double AUCt = 0.0;
		double AUCo = 0.0;
		double AUCi = 0.0;
		double AUCw = 0.0;
		double AUCs = 0.0;
		double AUCn = 0.0;
		meanISIs = new double[24];
		for(int i = 0; i < 24; i++){
			meanISIs[i] = (Math.random() * 2) + 4;
		}
		
		System.out.println("Enter the test number:");
		int TestNo = Keyboard.nextInt();
		System.out.println("Beginning the test.");
		
		long StartTime = System.nanoTime();
		
		for(int Z = 0; Z < 1; Z++){
			initializeNetwork();
			
			SetTrainingPattern GramGen = new SetTrainingPattern(corpFile);
			Corpus = GramGen.write(numTrainSentences);
			trainEpochs = Corpus.length;
			//D(n) correct outputs for this task
			double [][] Teach = _Matrix.GetSubArray(1, 0, trainEpochs, 24, Corpus);
			int [] corp = new int[Corpus.length];
			for(int i = 0, n = Corpus.length; i < n; i++) {
				for(int j = 0, m = Corpus[i].length; j < m; j++){
					if(Corpus[i][j] == 1.0){
						corp[i] = j;
						break;
					}
				}
			}		

			for(int i = 0; i < numRes; i++) {
				res[i].setReporter(false);
			}
			spikeGroups = createGroups();
			try {
				//**************************************************************//
				//*************************First Pass***************************//
				System.out.println("First-Pass...");
				latch = new CountDownLatch(1);		
				ExecutorService pool = Executors.newFixedThreadPool(numResThreads + 1);
				pool.submit(new runForward(corp, trainEpochs/2, stdp, null)).get();
				latch.await();
				//**************************************************************//
				//*************************TRAINING****************************//
				System.out.println("Training...");
				camera.readyFilm(trainEpochs, numRes);
				latch = new CountDownLatch(1);
				pool.submit(new runForward(corp, trainEpochs, null, camera)).get();
				latch.await();

				double [][] states = null;
				latch = new CountDownLatch(1);
				pool.submit(new solve(Teach, camera)).get();	
				latch.await();
				System.out.println("Network trained.");
	
				//**************************************************************//
				//*************************TESTING*****************************//
					 
				System.out.println("Testing...");
				SetTrainingPattern pattern =
						new SetTrainingPattern(corpFile, true);
				double [][] sents = pattern.write(numTestSentences);
				testEpochs = sents.length;
				
				double [][] targs = pattern.GetProbs().Get2DArray();
				pattern.GetProbs().Print("targs");
				double [][] outs;
				for(int i = 0, n = sents.length; i < n; i++) {
					for(int j = 0, m = sents[i].length; j < m; j++){
						if(sents[i][j] == 1.0){
							corp[i] = j;
							break;
						}
					}
				}
				camera.readyFilm(testEpochs, numRes);
				latch = new CountDownLatch(1);
				pool.submit(new runForward(corp, testEpochs, null, camera)).get();
				latch.await();
				pool.shutdown();
				double [] mean = new double[numRes + numInputs];
				states = Arrays.copyOfRange(camera.getCapturedStates(), 0, camera.getCapturedStates().length);
				for(int i = 0; i < states.length; i++) {
					states[i] = MathUtils.concatArrays(states[i], sents[i]);
				}
				Matrix S = new Matrix(states);
				Matrix tOuts = S.times(wOut);
				for(int i = 0; i < 24; i++) {
					for(int j = 0; j < tOuts.getRowDimension(); j++) {
						mean[i] += tOuts.get(j, i);
					}
					mean[i] = mean[i]/tOuts.getRowDimension();
				}	
				
				outs = tOuts.getArray();
				
				double [] SS = new double[outs[0].length];
				for(int i = 0; i < sents.length-1; i++){
					for(int h = 0; h < outs[0].length; h++){
						SS[h] += Math.pow(outs[i][h]-mean[h], 2);
						if(i == sents.length-2){
							SS[h] = Math.sqrt(SS[h]/(sents.length-2));
						}
					}
				}
				
				
				double [][] Zouts = new double[outs.length][outs[0].length];
				
				for(int i = 0; i < outs.length; i++){
					for(int J = 0; J < outs[0].length; J++){
						Zouts[i][J] = (outs[i][J] - mean[J])/SS[J];
						outs[i][J] += 0.5;
						if(outs[i][J] < 0.0){
							outs[i][J] = 0.0;
						}
					}
				}
				System.out.println("Outs: " + outs.length + " " + outs[0].length);
				System.out.println("Targs: " + targs.length + " " + targs[0].length);
			
				cos += Metrics.AvgCos(targs, outs);
				System.out.println("Avg Cosine Score: " + cos);
				MaxPred += Metrics.MaxPrediction(targs, Zouts);
				System.out.println("Max Prediction Score: " + MaxPred);
				VerbAgreement += Metrics.VerbAgreement(targs, Zouts);
				System.out.println("Verb Agreement Score: " + VerbAgreement);
				double [][] VDist = Metrics.VerbDist(targs, Zouts, 
						pattern.GetVerbDists());
				for(int i = 0; i < VerbD.length; i++){
					System.out.println(i+1 +": " + VerbD[i][0] + " ");
				}
				double [][] ROCc = Metrics.ROCurve(targs, Zouts);
				AUCn += Metrics.AUC(targs, Zouts, "Noun");
				AUCt += Metrics.AUC(targs, Zouts, "Trans");
				AUCo += Metrics.AUC(targs, Zouts, "Optrans");
				AUCi += Metrics.AUC(targs, Zouts, "Intrans");
				AUCw += Metrics.AUC(targs, Zouts, "Who");
				AUCs += Metrics.AUC(targs, Zouts, "Stop");
				
				for(int p = 0; p < VDist.length; p++){
					for(int q = 0; q < VDist[0].length; q++){
						VerbD[p][q] += VDist[p][q];
					}
				}
				for(int p = 0; p < ROCc.length; p++){
					for(int q = 0; q < ROCc[0].length; q++){
						ROC[p][q] += ROCc[p][q];
					}
			
				}
				
			} catch (Exception e){
				e.printStackTrace();
			}
			
			int Iter = Z+1;
			System.out.println("Iteration " + Iter + ": DONE.");
		}
		
		long End = System.nanoTime();

		double Sec = (double)(End-StartTime)/1000000000;
		System.out.println("Total time for 5 iterations: " + Sec + " sec.");
		
		try{
			FileWriter FW = new FileWriter("Metrics" + 
					numRes + "-"+TestNo+".txt");
			PrintWriter PW = new PrintWriter(FW);
			PW.println("Cos: " + cos);
			PW.println("MaxPred: " + MaxPred);
			PW.println("VerbAg: " + VerbAgreement);
			PW.println();
			PW.println("Verb Agreement Over Distance: ");
			for(int i = 0; i < VerbD.length; i++){
				PW.println(i+1 +": " + VerbD[i][0] +
					" Sample Size: " + VerbD[i][1]);
			}
			PW.println();
			PW.println("ROC area under the curve: ");
			PW.println("Nouns:              " + AUCn);
			PW.println("Transitive Verbs:   " + AUCt);
			PW.println("Intransitive Verbs: " + AUCi);
			PW.println("OpTransitive Verbs: " + AUCo);
			PW.println("Who:                " + AUCw);
			PW.println("Stop:               " + AUCs);
			PW.println();
			PW.println("Overall ROC curve: ");
			for(int i = 0; i < ROC.length; i++){
				for(int J = 0; J < ROC[0].length; J++){
					PW.print(ROC[i][J]+" ");
				}
				PW.println();
			}
			PW.close();

			
		}catch(IOException e){
			System.err.println(e);
		}
		
		System.exit(0);
	}
	
	public static HashMap<Integer, LIFNeuron[]> createGroups(){
		sgIndexes = new int [24][24];
		LIFNeuron [][] groups = new LIFNeuron[24][24];
		HashMap<Integer, LIFNeuron[]> groupFinal = new HashMap<Integer, LIFNeuron[]>();
		int c = 0;
		for(int i = 0; i < 24; i++){
			for(int j = 0; j<24;j++){
				sgIndexes[j][i] = c;
				groups[j][i] = inputs[c++];
			}	
		}
		
		for(int i = 0; i < 24; i++) {
			groupFinal.put(i, groups[i]);
		}
		
		
		return groupFinal;
	}
	
	
	
	public static void initializeNetwork() {
		
		camera = new StateCamera(CameraType.Bin,
				timeStep, exposureISI);
		
		gCon = new GridConnect();
		
		inputs = initializeInputs(numInputs);
		
		res = initializeNeurons(numRes, 30, 15, 5, 0.3, camera);
		rr = new Synapse [numRes][numRes];

		initializeSynVars(gCon);
			
		Network net = new Network();
		
		int [] inDims = {24, 24, 1};
		int [] resDims = {8, 8, 16};
		net.cubicLayout(new Coordinate(0,0,0), inDims , inputs);
		net.cubicLayout(new Coordinate(8, 8, 2), resDims, res);
		
		gCon.stochasticConnect(res, res, 3.0);
		//gCon.addPreferentialConnections(res);
		int y = 0;
		int i = 0;
		for(LIFNeuron n : res) {
			n.setIndex(i);
			i++;
		}
		i = 0;
		for(LIFNeuron n : res) {
			n.setIndex(y);
			for(Synapse s: n.getFanOut()) {
				rr[y][s.getTarg().getIndex()] = s;
				i++;
			}
			y++;
		}
		
		System.out.println("Res Sparsity: " + ((double) i)/(1024*1024));
		//TODO:Something important was here on this line then deleted. It probably doesn't matter though: I'm never using this class again.
		ir = gCon.stochasticConnect(inputs, res, 16);
		double lRate = 0.3 / Math.pow(10, 3);
		                // tcP,  tcD,  lRat,  nv,   alp,  bet, cp, wc 
		stdp = new LogSTDP(15,   30, lRate, 0.01f, 5.0f,  30f, 1f, 5.0f);
//		for(int q = 0; q < numRes; q++) {
//			res[q].transferSynapsesToArray();
//		}
//		for(int r = 0; r < numInputs; r++) {
//			inputs[r].transferSynapsesToArray();
//		}
	}	

	/** 
	 * Values from Maass (2002), D & F in mSecs
	 */
	public static void initializeSynVars (GridConnect gCon) {
		
		SynapseRandomizer eeSR = new SynapseRandomizer();
		eeSR.setU(new RandStream(Rand.GAUSS, 0.5, 0.25, 0.001,
				Double.MAX_VALUE, true));
		eeSR.setD(new RandStream(Rand.GAUSS, 1100, 550, 0.001,
				Double.MAX_VALUE, true));
		eeSR.setF(new RandStream(Rand.GAUSS, 50, 25, 0.001,
				Double.MAX_VALUE, true));
		
		
		SynapseRandomizer eiSR = new SynapseRandomizer();
		eiSR.setU(new RandStream(Rand.GAUSS, 0.05, 0.025, 0.001,
				Double.MAX_VALUE, true));
		eiSR.setD(new RandStream(Rand.GAUSS, 125, 62.5, 0.001,
				Double.MAX_VALUE, true));
		eiSR.setF(new RandStream(Rand.GAUSS, 120, 60, 0.001,
				Double.MAX_VALUE, true));
		
		
		SynapseRandomizer ieSR = new SynapseRandomizer();
		ieSR.setU(new RandStream(Rand.GAUSS, 0.25, 0.125, 0.001,
				Double.MAX_VALUE, true));
		ieSR.setD(new RandStream(Rand.GAUSS, 700, 350, 0.001,
				Double.MAX_VALUE, true));
		ieSR.setF(new RandStream(Rand.GAUSS, 20, 10, 0.001,
				Double.MAX_VALUE, true));
		
		
		SynapseRandomizer iiSR = new SynapseRandomizer();
		iiSR.setU(new RandStream(Rand.GAUSS, 0.32, 0.16, 0.0001,
				Double.MAX_VALUE, true));
		iiSR.setD(new RandStream(Rand.GAUSS, 144, 72, 0.001,
				Double.MAX_VALUE, true));
		iiSR.setF(new RandStream(Rand.GAUSS, 60, 30, 0.001,
				Double.MAX_VALUE, true));
		
		gCon.setEeSR(eeSR);
		gCon.setEiSR(eiSR);
		gCon.setIeSR(ieSR);
		gCon.setIiSR(iiSR);
		
	}
	
	public static LIFNeuron [] initializeNeurons(int numNeurons, int tc,
			double meanRate, double stdRate, double percentInhib,
				StateCamera camera) {
		LIFNeuron [] neurons = new LIFNeuron[numNeurons];	
	
		//double tc;
		for (int i = 0; i < numNeurons; i++) {
			//tc = RandStream.nextUniform(20, 40);
			if(Math.random() < percentInhib) {
				neurons[i] = new LIFNeuron(tc, false, i);
				neurons[i].setRefractory(2);
			} else {
				neurons[i] = new LIFNeuron(tc, true, i);
				neurons[i].setRefractory(3);
			}
			neurons[i].setPreferredISI(RandStream.nextGaussian(meanRate, stdRate, 5,
					meanRate*2, true));
			neurons[i].setBgCurrent(13.5);
			neurons[i].setHmP(true);
			if(camera != null) {
				neurons[i].setReporter(true);
				neurons[i].setRecorder(camera);
			}
		}
		
		initializeNeuronVars(neurons);
		return neurons;
	}
	
	public static LIFNeuron [] initializeInputs(int numInputs) {
		LIFNeuron [] neurons = new LIFNeuron[numInputs];	
		for (int i = 0; i < numInputs; i++) {
			neurons[i] = new LIFNeuron(30, true, i);
			neurons[i].setRefractory(0);
		}
		initializeNeuronVars(neurons);
		return neurons;
	}
	
	public static void initializeNeuronVars (LIFNeuron [] neurons) {
		
		for (int i = 0, n = neurons.length; i < n; i++) {
			neurons[i].setThreshold(15);
			neurons[i].setLastSpike(0);
			neurons[i].setTime(0);
			neurons[i].setReset(13.5);
			neurons[i].setMemResistance(1);
			neurons[i].setRestPotential(0);	
			neurons[i].setMemPotential(RandStream.nextUniform(13.5, 15.1));	
		}		
		
	}
}


